import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib的字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体为微软雅黑
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题

# 参数定义
N = 10000  # 总人口
beta = 0.2  # 感染率
gamma = 0.1  # 疫苗接种率
delta = 0.1  # 暴露者转化为感染者的速率
alpha = 0.05  # 感染者恢复率
rho = 0.01  # 感染者死亡率
dt = 0.1  # 时间步长
T = 100  # 总时间

# 初始化
S0, V0, E0, I0, R0, D0 = N - 1, 0, 0, 1, 0, 0  # 初始条件
np.random.seed(0)


# 函数：执行Euler-Maruyama模拟并返回结果
def simulate(sigma):
    S, V, E, I, R, D = [S0], [V0], [E0], [I0], [R0], [D0]
    for t in np.arange(dt, T, dt):
        dW_S, dW_V, dW_E, dW_I, dW_R, dW_D = np.random.normal(0, np.sqrt(dt), 6)
        S_new = S[-1] + (-beta * S[-1] * (I[-1] + E[-1]) / N - gamma * S[-1]) * dt + sigma * S[-1] * dW_S
        V_new = V[-1] + (gamma * S[-1] - sigma * V[-1] * (I[-1] + E[-1]) / N) * dt + sigma * V[-1] * dW_V
        E_new = E[-1] + (
                beta * S[-1] * (I[-1] + E[-1]) / N + sigma * V[-1] * (I[-1] + E[-1]) / N - delta * E[-1]) * dt + sigma * \
                E[-1] * dW_E
        I_new = I[-1] + (delta * E[-1] - (alpha + rho) * I[-1]) * dt + sigma * I[-1] * dW_I
        R_new = R[-1] + alpha * I[-1] * dt + sigma * R[-1] * dW_R
        D_new = D[-1] + rho * I[-1] * dt + sigma * D[-1] * dW_D
        S.append(S_new)
        V.append(V_new)
        E.append(E_new)
        I.append(I_new)
        R.append(R_new)
        D.append(D_new)
    return S, V, E, I, R, D


# 波动幅度分析图
sigmas = [0.01, 0.05, 0.1]  # 不同的波动强度
colors = ['blue', 'orange', 'green']  # 对应曲线的颜色

# 为每个sigma值执行并保存结果
for i, sigma in enumerate(sigmas):
    S, V, E, I, R, D = simulate(sigma)
    plt.plot(np.arange(0, T, dt), I, label=f'σ={sigma} 传染者', color=colors[i])

# 绘制图形
plt.xlabel('时间')
plt.ylabel('传染者人数')
plt.title('不同随机波动强度下的疾病传播情况')
plt.legend()
plt.grid(True)
plt.show()